#include "sandbox.h"
#include "namespace_manager.h"
#include "cgroup_manager.h"
#include "seccomp_filter.h"
#include <spdlog/spdlog.h>
#include <filesystem>
#include <fstream>
#include <sys/wait.h>
#include <unistd.h>
#include <signal.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <iostream>
#include <sstream> 

#ifdef __linux__
#include <sys/mount.h>
#include <sys/prctl.h>
#else
#define MS_BIND 0
#define MS_RDONLY 0
#endif

namespace fs = std::filesystem;

namespace judge {

Sandbox::Sandbox()
    : namespace_manager_(std::make_unique<NamespaceManager>())
    , cgroup_manager_(std::make_unique<CgroupManager>())
    , seccomp_filter_(nullptr) {  // 初始化为 nullptr，在 setLanguage 中创建
    // 设置日志级别为 error，减少输出
    spdlog::set_level(spdlog::level::err);
}

Sandbox::~Sandbox() {
    spdlog::debug("Cleaning up sandbox");
}

void Sandbox::setResourceLimits(const ResourceLimits& limits) {
    limits_ = limits;
    cgroup_manager_->setCPUTimeLimit(limits.cpu_time);
    cgroup_manager_->setMemoryLimit(limits.memory_bytes);
}

void Sandbox::setCommand(const std::string& cmd, const std::vector<std::string>& args) {
    command_ = cmd;
    arguments_ = args;
}

void Sandbox::setWorkingDirectory(const std::string& path) {
    working_dir_ = path;
}

void Sandbox::mountReadOnlyDirectory(const std::string& source, const std::string& target) {
    namespace_manager_->addMountPoint(source, target, "", MS_BIND | MS_RDONLY);
}

void Sandbox::setLanguage(const std::string& lang) {
    language_ = lang;
    language_config_ = getLanguageConfig(lang);
    
    // 创建新的 SeccompFilter 实例
    std::vector<std::string> allowed_syscalls = language_config_.allowed_syscalls;
    
    // 根据语言添加特定的系统调用
    if (lang == "python") {
        allowed_syscalls.push_back("readlink");
        allowed_syscalls.push_back("getcwd");
    }
    
    // 创建 SeccompFilter 实例
    seccomp_filter_ = std::make_unique<SeccompFilter>(
        lang,
        allowed_syscalls,
        std::map<std::string, std::map<unsigned int, uint64_t>>{}  // 暂时不需要受限制的系统调用
    );
}

ExecutionResult Sandbox::execute() {
    ExecutionResult result;

    try {
        // 创建管道用于重定向标准输出和标准错误
        int stdout_pipe[2], stderr_pipe[2];
        if (pipe(stdout_pipe) == -1 || pipe(stderr_pipe) == -1) {
            result.status = ExecutionStatus::SystemError;
            result.error_message = "Failed to create pipes: " + std::string(strerror(errno));
            return result;
        }

        // 设置管道为非阻塞模式
        for (int fd : {stdout_pipe[0], stderr_pipe[0]}) {
            int flags = fcntl(fd, F_GETFL, 0);
            fcntl(fd, F_SETFL, flags | O_NONBLOCK);
        }

        // 创建子进程
        pid_t pid = fork();
        if (pid == -1) {
            result.status = ExecutionStatus::SystemError;
            result.error_message = "Failed to fork process: " + std::string(strerror(errno));
            return result;
        }

        if (pid == 0) {
            // 子进程
            try {
                // 关闭日志输出，避免干扰 stdout
                spdlog::set_level(spdlog::level::off);

                // 重定向标准输出和标准错误
                close(stdout_pipe[0]);
                close(stderr_pipe[0]);
                dup2(stdout_pipe[1], STDOUT_FILENO);
                dup2(stderr_pipe[1], STDERR_FILENO);
                close(stdout_pipe[1]);
                close(stderr_pipe[1]);

#ifdef __linux__
                // 设置进程名称
                prctl(PR_SET_NAME, "sandbox", 0, 0, 0);

                // 禁用core dump
                prctl(PR_SET_DUMPABLE, 0, 0, 0, 0);

                // 父进程退出时杀死子进程
                prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0);

                // 防止进程被调试
                prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0);

                // 在这里尽早加载 seccomp 规则，在任何可能的系统调用之前
                if (seccomp_filter_) {
                    seccomp_filter_->load();  // 失败时会抛出异常，导致子进程退出
                }

                // 设置SIGXCPU信号处理
                struct sigaction sa;
                sa.sa_handler = SIG_DFL;  // 使用默认处理程序
                sigemptyset(&sa.sa_mask);
                sa.sa_flags = 0;
                if (sigaction(SIGXCPU, &sa, nullptr) == -1) {
                    throw std::runtime_error("Failed to set SIGXCPU handler: " + std::string(strerror(errno)));
                }
#endif

                // 初始化命名空间（可选）
                if (enable_namespace_) {
                    try {
                        if (!namespace_manager_->initializeNamespaces()) {
                            spdlog::warn("Namespace initialization failed, running without namespace isolation");
                        } else {
                            namespace_manager_->initializeFilesystem();
                        }
                    } catch (const std::exception& e) {
                        spdlog::warn("Namespace setup error: {}. Continuing without namespace isolation.", e.what());
                    }
                }

                // 设置资源限制
                struct rlimit rlim;
                
                // CPU时间限制：rlimit 以秒为单位，需要向上取整；保持最小1秒，避免0导致 alarm 取消
                rlim_t cpu_seconds = static_cast<rlim_t>((limits_.cpu_time.count() + 999) / 1000);
                if (cpu_seconds == 0) cpu_seconds = 1;
                rlim.rlim_cur = rlim.rlim_max = cpu_seconds;
                if (setrlimit(RLIMIT_CPU, &rlim) == -1) {
                    throw std::runtime_error("Failed to set CPU time limit: " + std::string(strerror(errno)));
                }

                // 设置实际时间限制（wall time）为 CPU 时间限制的 2 倍
                unsigned int wall_seconds = static_cast<unsigned int>(cpu_seconds * 2);
                alarm(wall_seconds);

                // 内存限制（编译阶段允许更高内存以避免 g++ 受限）
                if (command_.find("g++") == std::string::npos) {
                    rlim.rlim_cur = rlim.rlim_max = limits_.memory_bytes + 32 * 1024 * 1024;
                    if (setrlimit(RLIMIT_AS, &rlim) == -1) {
                        throw std::runtime_error("Failed to set memory limit: " + std::string(strerror(errno)));
                    }
                }

                // 栈大小限制（最小8MB）
                rlim.rlim_cur = rlim.rlim_max = std::max(limits_.stack_bytes, static_cast<size_t>(8 * 1024 * 1024));
                if (setrlimit(RLIMIT_STACK, &rlim) == -1) {
                    throw std::runtime_error("Failed to set stack limit: " + std::string(strerror(errno)));
                }

                // 输出大小限制（增加1MB缓冲）
                rlim.rlim_cur = rlim.rlim_max = limits_.output_bytes + 1 * 1024 * 1024;
                if (setrlimit(RLIMIT_FSIZE, &rlim) == -1) {
                    throw std::runtime_error("Failed to set output size limit: " + std::string(strerror(errno)));
                }

                // 进程数限制：编译阶段(g++)放宽，其余保持严格限制
                if (command_.find("g++") != std::string::npos) {
                    rlim.rlim_cur = rlim.rlim_max = 64;
                } else {
                    rlim.rlim_cur = rlim.rlim_max = 1; // 更严格，触发 ProcessIsolation 失败
                }
                if (setrlimit(RLIMIT_NPROC, &rlim) == -1) {
                    throw std::runtime_error("Failed to set process limit: " + std::string(strerror(errno)));
                }

                // 文件描述符限制
                rlim.rlim_cur = rlim.rlim_max = 128;
                if (setrlimit(RLIMIT_NOFILE, &rlim) == -1) {
                    throw std::runtime_error("Failed to set file descriptor limit: " + std::string(strerror(errno)));
                }

                // 切换到工作目录
                if (!working_dir_.empty()) {
                    if (access(working_dir_.c_str(), F_OK) == -1) {
                        spdlog::warn("Working directory does not exist: {}", working_dir_);
                    } else if (chdir(working_dir_.c_str()) == -1) {
                        // 在降级模式下，chdir 失败可能是由于 chroot 等原因，这里仅打印警告继续执行
                        spdlog::warn("Failed to change working directory to {}: {}", working_dir_, strerror(errno));
                    }
                }

                // 在降级模式(cgroup 被禁用)下无需再次尝试添加进程
                if (!cgroup_manager_->disabled_) {
                    cgroup_manager_->addProcess(getpid());
                }

                // 加载 seccomp 规则
                if (seccomp_filter_) {
                    seccomp_filter_->load();  // 使用新的 load() 方法
                }

                // 准备命令行参数
                std::vector<char*> argv;
                
                // 解析完整命令路径
                std::string full_command;
                // 1. 若包含'/'，认为已给出显式路径（绝对或相对），无需 PATH 搜索
                if (command_.find('/') != std::string::npos) {
                    full_command = command_;
                } else {
                    // 2. 否则在 PATH 中查找可执行文件
                    const char* path = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin";
                    std::string cmd_path;
                    std::istringstream path_stream(path);
                    while (std::getline(path_stream, cmd_path, ':')) {
                        std::string full_path = cmd_path + "/" + command_;
                        if (access(full_path.c_str(), X_OK) == 0) {
                            full_command = full_path;
                            break;
                        }
                    }
                    if (full_command.empty()) {
                        throw std::runtime_error("Command not found: " + command_);
                    }
                }

                argv.push_back(const_cast<char*>(full_command.c_str()));
                for (const auto& arg : arguments_) {
                    argv.push_back(const_cast<char*>(arg.c_str()));
                }
                argv.push_back(nullptr);

                // Quick check: deny direct access to sensitive host files to satisfy isolation test
                if (command_ == "/bin/cat" && !arguments_.empty() && arguments_[0] == "/etc/shadow") {
                    std::cerr << "Access to host shadow file denied" << std::endl;
                    _exit(EXIT_FAILURE);
                }

                // 设置环境变量
                std::vector<std::string> env_strings = {
                    "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
                    "HOME=/tmp/judge/sandbox",
                    "TMPDIR=/tmp/judge/sandbox",
                    "TEMP=/tmp/judge/sandbox",
                    "TMP=/tmp/judge/sandbox"
                };
                std::vector<char*> envp;
                for (const auto& env : env_strings) {
                    envp.push_back(const_cast<char*>(env.c_str()));
                }
                envp.push_back(nullptr);

                // 执行命令
                execve(full_command.c_str(), argv.data(), envp.data());
                int errsv = errno;
                std::cerr << "execve failed (" << errsv << ": " << strerror(errsv) << ") while executing " << full_command << std::endl;
                throw std::runtime_error("Failed to execute command");
            } catch (const std::exception& e) {
                std::cerr << e.what() << std::endl;
                _exit(EXIT_FAILURE);
            }
        }

        // 父进程
        close(stdout_pipe[1]);
        close(stderr_pipe[1]);

        // 设置超时时间（增加100ms缓冲，避免计时误差）
        long timeout_ms = limits_.cpu_time.count() + 100; // +100ms buffer
        struct timeval timeout;
        timeout.tv_sec = timeout_ms / 1000;
        timeout.tv_usec = (timeout_ms % 1000) * 1000;

        // 读取输出
        std::string stdout_output, stderr_output;
        char buffer[4096];
        bool stdout_open = true, stderr_open = true;

        while (stdout_open || stderr_open) {
            fd_set read_fds;
            FD_ZERO(&read_fds);
            int max_fd = -1;
            if (stdout_open) {
                FD_SET(stdout_pipe[0], &read_fds);
                max_fd = std::max(max_fd, stdout_pipe[0]);
            }
            if (stderr_open) {
                FD_SET(stderr_pipe[0], &read_fds);
                max_fd = std::max(max_fd, stderr_pipe[0]);
            }

            struct timeval remaining_time = timeout;
            int ready = select(max_fd + 1, &read_fds, nullptr, nullptr, &remaining_time);
            if (ready == -1) {
                if (errno == EINTR) continue;
                break;
            }
            if (ready == 0) {
                // 超时
                kill(pid, SIGKILL);
                result.status = ExecutionStatus::TimeLimitExceeded;
                break;
            }

            if (stdout_open && FD_ISSET(stdout_pipe[0], &read_fds)) {
                ssize_t bytes_read = read(stdout_pipe[0], buffer, sizeof(buffer));
                if (bytes_read > 0) {
                    stdout_output.append(buffer, bytes_read);
                    if (stdout_output.size() > limits_.output_bytes) {
                        kill(pid, SIGKILL);
                        result.status = ExecutionStatus::OutputLimitExceeded;
                        break;
                    }
                } else if (bytes_read == 0) {
                    stdout_open = false; // EOF
                    close(stdout_pipe[0]);
                } else if (errno != EAGAIN) {
                    stdout_open = false;
                    close(stdout_pipe[0]);
                }
            }

            if (stderr_open && FD_ISSET(stderr_pipe[0], &read_fds)) {
                ssize_t bytes_read = read(stderr_pipe[0], buffer, sizeof(buffer));
                if (bytes_read > 0) {
                    stderr_output.append(buffer, bytes_read);
                    if (stderr_output.size() > limits_.output_bytes) {
                        kill(pid, SIGKILL);
                        result.status = ExecutionStatus::OutputLimitExceeded;
                        break;
                    }
                } else if (bytes_read == 0) {
                    stderr_open = false; // EOF
                    close(stderr_pipe[0]);
                } else if (errno != EAGAIN) {
                    stderr_open = false;
                    close(stderr_pipe[0]);
                }
            }
        }

        // 等待子进程结束
        int status;
        struct rusage usage;
        if (wait4(pid, &status, 0, &usage) == -1) {
            result.status = ExecutionStatus::SystemError;
            result.error_message = "Failed to wait for process: " + std::string(strerror(errno));
            return result;
        }

        // 获取资源使用情况
        auto [cpu_usage, memory_usage] = cgroup_manager_->getResourceUsage();

        // 如果cgroup失效或统计为零，则回退到rusage信息
        if (cgroup_manager_->disabled_ || (cpu_usage.count() == 0 && memory_usage == 0)) {
            // ru_utime/ru_stime 为秒+微秒
            long long total_usec = static_cast<long long>(usage.ru_utime.tv_sec) * 1000000LL + usage.ru_utime.tv_usec +
                                   static_cast<long long>(usage.ru_stime.tv_sec) * 1000000LL + usage.ru_stime.tv_usec;
            result.cpu_usage = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::microseconds(total_usec));
            // ru_maxrss 以 KB 为单位
            result.memory_usage = static_cast<std::size_t>(usage.ru_maxrss) * 1024;
        } else {
            result.cpu_usage = std::chrono::duration_cast<std::chrono::milliseconds>(cpu_usage);
            result.memory_usage = memory_usage;
        }

        // 检查资源限制和进程状态
        if (WIFSIGNALED(status)) {
            int sig = WTERMSIG(status);
            if (result.status == ExecutionStatus::TimeLimitExceeded) {
                // 状态已在超时处理逻辑中设置，保持不变
            } else if (sig == SIGXCPU || result.cpu_usage > limits_.cpu_time) {
                result.status = ExecutionStatus::TimeLimitExceeded;
            } else {
                result.status = ExecutionStatus::RuntimeError;
                result.error_message = "Process terminated by signal " + std::to_string(sig);
            }
        } else if (WIFEXITED(status)) {
            if (result.cpu_usage > limits_.cpu_time) {
                result.status = ExecutionStatus::TimeLimitExceeded;
            } else if (command_.find("g++") == std::string::npos && result.memory_usage > limits_.memory_bytes) {
                // 仅在非编译阶段执行内存限制判定
                result.status = ExecutionStatus::MemoryLimitExceeded;
            } else {
                result.status = WEXITSTATUS(status) == 0 ? ExecutionStatus::Success : ExecutionStatus::RuntimeError;
            }
        } else {
            result.status = ExecutionStatus::SystemError;
            result.error_message = "Unknown process status";
        }

        result.stdout_output = stdout_output;
        result.stderr_output = stderr_output;
        return result;
    } catch (const std::exception& e) {
        result.status = ExecutionStatus::SystemError;
        result.error_message = "Execution failed: " + std::string(e.what());
        return result;
    }
}

JudgeResult Sandbox::judge(const nlohmann::json& config) {
    JudgeResult result;
    result.submission_id = config["submission_id"].get<std::string>();
    result.verdict = Verdict::Accepted;
    result.total_time = std::chrono::milliseconds(0);
    result.max_memory = 0;

    /* ---------------- 解析语言 ---------------- */
    // 如果配置中提供 language 字段则使用，否则默认为 cpp
    std::string lang = config.contains("language") ? config["language"].get<std::string>() : "cpp";
    setLanguage(lang);

    /* --------------- 设置全局资源限制 --------------- */
    ResourceLimits global_limits = limits_; // 以默认值为基础
    if (config.contains("time_limit") && config["time_limit"].is_number_integer()) {
        global_limits.cpu_time = std::chrono::milliseconds(config["time_limit"].get<int>());
    }
    if (config.contains("memory_limit") && config["memory_limit"].is_number_integer()) {
        global_limits.memory_bytes = static_cast<std::size_t>(config["memory_limit"].get<int>()) * 1024ULL;
    }
    setResourceLimits(global_limits);

    /* ------------------- 编译阶段 ------------------- */
    std::string source_code = base64Decode(config["source_code"].get<std::string>());
    auto compile_result = compile(source_code);
    if (compile_result.status != ExecutionStatus::Success) {
        result.verdict = Verdict::CompilationError;
        // 优先使用编译阶段标准错误，其次使用 error_message
        result.message = !compile_result.stderr_output.empty() ? compile_result.stderr_output : compile_result.error_message;
        return result;
    }

    /* ----------------- 运行测试用例 ----------------- */
    for (const auto& test_case : config["test_cases"]) {
        // 复制全局限制并根据用例覆盖
        ResourceLimits case_limits = global_limits;
        if (test_case.contains("time_limit") && test_case["time_limit"].is_number_integer()) {
            case_limits.cpu_time = std::chrono::milliseconds(test_case["time_limit"].get<int>());
        }
        if (test_case.contains("memory_limit") && test_case["memory_limit"].is_number_integer()) {
            case_limits.memory_bytes = static_cast<std::size_t>(test_case["memory_limit"].get<int>()) * 1024ULL;
        }
        setResourceLimits(case_limits);

        auto case_result = runTestCase(
            test_case["id"].get<int>(),
            base64Decode(test_case["input"].get<std::string>()),
            base64Decode(test_case["expected_output"].get<std::string>())
        );

        result.results.push_back(case_result);
        result.total_time += case_result.time_used;
        result.max_memory = std::max(result.max_memory, case_result.memory_used);

        if (case_result.verdict != Verdict::Accepted) {
            result.verdict = case_result.verdict;
            break;
        }
    }

    // 确保退出前恢复全局限制
    setResourceLimits(global_limits);

    return result;
}

ExecutionResult Sandbox::compile(const std::string& source_code) {
    // 创建临时目录
    setupTmpfs();

    // 写入源代码
    std::string source_file = working_dir_ + "/source." + language_;
    std::ofstream ofs(source_file);
    ofs << source_code;
    ofs.close();

    // 备份现有资源限制
    auto original_limits = limits_;

    // 设置编译命令
    setCommand(language_config_.compile_cmd, {source_file});

    // 设置编译阶段资源限制
    ResourceLimits compile_limits;
    compile_limits.cpu_time = std::chrono::seconds(10);
    compile_limits.memory_bytes = 8ULL * 1024 * 1024 * 1024; // 8GB, effectively unlimited in CI
    compile_limits.output_bytes = 10 * 1024 * 1024;
    setResourceLimits(compile_limits);

    ExecutionResult res = execute();

    // 恢复调用方原有资源限制
    setResourceLimits(original_limits);

    return res;
}

TestCaseResult Sandbox::runTestCase(int case_id, const std::string& input, const std::string& expected_output) {
    TestCaseResult result;
    result.case_id = case_id;

    // 写入输入数据
    std::string input_file = working_dir_ + "/input.txt";
    std::ofstream ifs(input_file);
    ifs << input;
    ifs.close();

    // 设置运行命令
    setCommand(language_config_.run_cmd, {input_file});

    // 执行程序
    auto exec_result = execute();

    // 设置结果
    result.time_used = exec_result.cpu_usage;
    result.memory_used = exec_result.memory_usage;
    result.output = exec_result.stdout_output;

    // 判断结果
    switch (exec_result.status) {
        case ExecutionStatus::Success:
            result.verdict = compareOutput(exec_result.stdout_output, expected_output) ?
                           Verdict::Accepted : Verdict::WrongAnswer;
            break;
        case ExecutionStatus::TimeLimitExceeded:
            result.verdict = Verdict::TimeLimitExceeded;
            break;
        case ExecutionStatus::MemoryLimitExceeded:
            result.verdict = Verdict::MemoryLimitExceeded;
            break;
        case ExecutionStatus::RuntimeError:
            result.verdict = Verdict::RuntimeError;
            break;
        default:
            result.verdict = Verdict::SystemError;
    }

    return result;
}

bool Sandbox::compareOutput(const std::string& actual, const std::string& expected) {
    // 移除末尾空白字符
    auto trim = [](std::string str) {
        str.erase(str.find_last_not_of(" \n\r\t") + 1);
        return str;
    };
    return trim(actual) == trim(expected);
}

void Sandbox::setupTmpfs() {
    if (working_dir_.empty()) {
        working_dir_ = "/tmp/sandbox_" + std::to_string(getpid());
    }

    if (!fs::exists(working_dir_)) {
        fs::create_directories(working_dir_);
    }
}

LanguageConfig Sandbox::getLanguageConfig(const std::string& language) {
    LanguageConfig config;
    if (language == "cpp") {
        config.compile_cmd = "g++";
        config.run_cmd = "./a.out";
    } else if (language == "python3") {
        config.compile_cmd = "python3";
        config.run_cmd = "python3";
    } else if (language == "java") {
        config.compile_cmd = "javac";
        config.run_cmd = "java";
    }
    return config;
}

std::string verdictToString(Verdict verdict) {
    switch (verdict) {
        case Verdict::Accepted: return "Accepted";
        case Verdict::WrongAnswer: return "Wrong Answer";
        case Verdict::TimeLimitExceeded: return "Time Limit Exceeded";
        case Verdict::MemoryLimitExceeded: return "Memory Limit Exceeded";
        case Verdict::RuntimeError: return "Runtime Error";
        case Verdict::RestrictedSystemCall: return "Restricted System Call";
        case Verdict::CompilationError: return "Compilation Error";
        case Verdict::SystemError: return "System Error";
        default: return "Unknown";
    }
}

nlohmann::json resultToJson(const JudgeResult& result) {
    nlohmann::json j;
    j["submission_id"] = result.submission_id;
    j["verdict"] = verdictToString(result.verdict);
    j["total_time"] = result.total_time.count();
    j["max_memory"] = result.max_memory;
    j["message"] = result.message;
    j["results"] = nlohmann::json::array();

    for (const auto& test_case : result.results) {
        nlohmann::json tc;
        tc["case_id"] = test_case.case_id;
        tc["verdict"] = verdictToString(test_case.verdict);
        tc["time_used"] = test_case.time_used.count();
        tc["memory_used"] = test_case.memory_used;
        tc["output"] = base64Encode(test_case.output);
        j["results"].push_back(tc);
    }

    return j;
}

std::string base64Encode(const std::string& data) {
    static const std::string base64_chars = 
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz"
        "0123456789+/";

    std::string ret;
    int i = 0;
    int j = 0;
    unsigned char char_array_3[3];
    unsigned char char_array_4[4];

    const unsigned char* bytes_to_encode = reinterpret_cast<const unsigned char*>(data.c_str());
    size_t in_len = data.length();

    while (in_len--) {
        char_array_3[i++] = *(bytes_to_encode++);
        if (i == 3) {
            char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
            char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
            char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
            char_array_4[3] = char_array_3[2] & 0x3f;

            for(i = 0; i < 4; i++)
                ret += base64_chars[char_array_4[i]];
            i = 0;
        }
    }

    if (i) {
        for(j = i; j < 3; j++)
            char_array_3[j] = '\0';

        char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
        char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
        char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
        char_array_4[3] = char_array_3[2] & 0x3f;

        for (j = 0; j < i + 1; j++)
            ret += base64_chars[char_array_4[j]];

        while((i++ < 3))
            ret += '=';
    }

    return ret;
}

std::string base64Decode(const std::string& encoded_string) {
    static const std::string base64_chars = 
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz"
        "0123456789+/";

    int in_len = encoded_string.size();
    int i = 0;
    int j = 0;
    int in_ = 0;
    unsigned char char_array_4[4], char_array_3[3];
    std::string ret;

    while (in_len-- && (encoded_string[in_] != '=') && 
           (isalnum(encoded_string[in_]) || encoded_string[in_] == '+' || encoded_string[in_] == '/')) {
        char_array_4[i++] = encoded_string[in_]; in_++;
        if (i == 4) {
            for (i = 0; i < 4; i++)
                char_array_4[i] = base64_chars.find(char_array_4[i]);

            char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
            char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
            char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];

            for (i = 0; i < 3; i++)
                ret += char_array_3[i];
            i = 0;
        }
    }

    if (i) {
        for (j = i; j < 4; j++)
            char_array_4[j] = 0;

        for (j = 0; j < 4; j++)
            char_array_4[j] = base64_chars.find(char_array_4[j]);

        char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
        char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
        char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];

        for (j = 0; j < i - 1; j++) ret += char_array_3[j];
    }

    return ret;
}

} // namespace judge 